# -*- coding: utf-8 -*-
# @Time    : 2023/6/2 13:47
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: ClasBaseDataset

import os
import paddle
from paddle import io
from datasets import base
from transforms import Compose


class ClasBaseDataset(io.Dataset):
    def __init__(self, config):
        super(ClasBaseDataset, self).__init__()
        self.mode = config["mode"] if "mode" in config.keys() else "standard"
        self.data_root = config["data_root"] if "data_root" in config.keys() else None
        self.data_list_file = config["data_list"] if "data_list" in config.keys() else None
        self.split_sign = config["split_sign"] if "split_sign" in config.keys() else None
        self.trans = Compose(config["transforms"])
        self.data_list = self._make_list()

    def __getitem__(self, item):
        data = eval("self._"+self.mode)(item)
        return data

    def __len__(self):
        return len(self.data_list)

    def _standard(self, item):
        data = {"img": os.path.join(self.data_root, self.data_list[item][0]),
                "path": self.data_list[item],
                "label": self.data_list[item][1]}
        data = self.trans(data)
        return data

    def _predict(self, item):
        data = {"img": os.path.join(self.data_root, self.data_list[item]),
                "path": self.data_list[item]}
        data = self.trans(data)
        return data

    def _make_list(self):
        if self.mode:
            assert self.data_list_file is not None
            data_list_file = open(self.data_list_file, "r")
            data_list = data_list_file.read().split("\n")[:]
            data_list_file.close()
            data_list = [[data.split()[0], int(data.split()[1])] for data in data_list]
        else:
            data_list = base.recursion(self.data_root, self.recursion_identifier)
        return data_list
